import { Inject, Injectable } from '@nestjs/common';
import { RecursiveCharacterTextSplitter } from 'langchain/text_splitter';
import { Document } from '@langchain/core/documents';
import { WeaviateStore } from '@langchain/weaviate';
import { Filters, type WeaviateClient } from 'weaviate-client';

import { EmbeddingsService } from '../embeddings/embeddings.service';

@Injectable()
export class WeaviateService {
  textSplitter: RecursiveCharacterTextSplitter;
  vectorStore: WeaviateStore;

  constructor(
    readonly embeddingsService: EmbeddingsService,
    @Inject('WEAVIATE_CLIENT')
    private readonly weaviateClient: WeaviateClient,
  ) {
    this.textSplitter = new RecursiveCharacterTextSplitter({
      chunkSize: 500,
      chunkOverlap: 50,
    });

    this.vectorStore = new WeaviateStore(embeddingsService.embeddings, {
      client: this.weaviateClient as any,
      indexName: 'Langchainjs_test',
    });
  }

  async addDocuments(docs: Document[]) {
    const splitDocs = await this.textSplitter.splitDocuments(docs);

    return await this.vectorStore.addDocuments(splitDocs);
  }

  async findAll(current: number = 1, pageSize: number = 10, source: string) {
    const collection = this.weaviateClient.collections.use('Langchainjs_test');

    const filters = Filters.and(collection.filter.byProperty('source').equal(source));

    const result = await collection.query.fetchObjects({
      limit: pageSize,
      offset: (current - 1) * pageSize,
      filters: source ? filters : undefined,
      returnMetadata: 'all',
    });

    const datas = await collection.aggregate.overAll();

    return {
      data: result.objects,
      total: datas.totalCount,
      current,
      pageSize,
      pages: Math.ceil(datas.totalCount / pageSize),
    };
  }

  async queryDocuments(query: string, k: number) {
    return this.vectorStore.similaritySearchWithScore(query, k);
  }

  async deleteDocuments() {
    await this.weaviateClient.collections.delete('Langchainjs_test');
    await this.weaviateClient.collections.create({
      name: 'Langchainjs_test',
      properties: [
        { name: 'source', dataType: 'text' },
        { name: 'content', dataType: 'text' },
      ],
    });
  }
}
